# -*- coding: utf-8 -*-
"""
Created on Sun Oct 19 10:45:30 2025

@author: xiaok
"""


import mne
import os
import math
import random
import numpy as np
import pandas as pd
import seaborn as sn
import matplotlib.pyplot as plt
from scipy.stats import f_oneway
from sklearn.preprocessing import StandardScaler


raw_data_folder = 'D:/materials/dataset/bci/BCICIV_2a_gdf/'

files = os.listdir(raw_data_folder)

plt.close('all')


#%% 
# List of the events
# '1023': 1 Rejected trial
# '1072': 2 Eye movements
# '276': 3 Idling EEG (eyes open)
# '277': 4 Idling EEG (eyes closed)
# '32766': 5 Start of a new run
# '768': 6 Start of a trial
# '769': 7 Cue onset Left (class 1) : 0
# '770': 8 Cue onset Right (class 2) : 1
# '771': 9 Cue onset Foot (class 3) : 2
# '772': 10 Cue onset Tongue (class 4): 3

labels = np.empty([0], dtype = int) 
tmin = 0
tmax = 4

def StandardScalar_data(data):
    print(data.shape)
    # should be (288, 22, 1001)
    
    n_samples, n_channels, n_timepoints = data.shape
    data_flat = data.reshape(n_samples,-1)
    print(data_flat.shape)
    # should be (288, 22022)

    scaler = StandardScaler().fit(data_flat)
    data_scaled = scaler.transform(data_flat)
    
    data_scaled = data_scaled.reshape(n_samples, n_channels, n_timepoints)
    print(data_scaled.shape)
    # should be (288, 22, 1001)

    return data_scaled

for i in range(1,9):
    fileName = 'A0'+str(i)+'T.gdf'
    print(fileName)
    file_path = os.path.join(raw_data_folder, fileName)
    raw = mne.io.read_raw_gdf(file_path, eog=['EOG-left', 'EOG-central', 'EOG-right'], preload=True)
    raw.drop_channels(['EOG-left', 'EOG-central', 'EOG-right'])
    
    events_from_annot,event_dict = mne.events_from_annotations(raw)
    event_id = [event_dict['769'],event_dict['770'],event_dict['771'],event_dict['772']]
    # fig = mne.viz.plot_events(events_from_annot,event_id=event_dict,sfreq=raw.info['sfreq'],first_samp=raw.first_samp)
    # raw.resample(128, npad='auto')
    
    #     # High Pass Filtering 4-40 Hz
    #     raw.filter(l_freq=1, h_freq=100, method='iir')

    #     # Notch filter for Removal of Line Voltage
    #     raw.notch_filter(freqs=50)
    epochs = mne.Epochs(raw, events_from_annot, event_id=event_id, tmin=tmin, tmax=tmax, reject=None,proj=False, baseline=None, preload=True)
    tmp_data =  epochs.get_data()
    tmp_data_2 = StandardScalar_data(tmp_data)    
    
    if i==1:
        data = tmp_data_2        
    else:
        data = np.concatenate((data, tmp_data_2), axis=0)
        
        
    
    # convert the event_ids(6,7 ..) to original labels like 769,770
    labels_tmp = epochs.events[:,-1]
    labels_int = []
    for l in labels_tmp:
        keys = [key for key, val in event_dict.items() if val == l]
        labels_int.append(int(''.join(keys)))
    labels = np.concatenate((labels,labels_int), axis=0)

#%%

import pickle

bcic_iv_2a_data_all_sub = {}
bcic_iv_2a_data_all_sub['datas']=data
bcic_iv_2a_data_all_sub['labels']=labels
bcic_iv_2a_data_all_sub['fs']=raw.info['sfreq']

#%%


filePathName_save = 'D:/materials/dataset/bci/BCICIV_2a_pickle/bcic_iv_2a_data_T.p'
pickle.dump(bcic_iv_2a_data_all_sub,open(filePathName_save,'wb'))





